package task

import (
	"context"
	"fmt"
	"sync"
	"time"
)

// Task 任务接口
type Task interface {
	// Execute 执行任务
	Execute(ctx context.Context) error
	// GetID 获取任务ID
	GetID() string
	// GetType 获取任务类型
	GetType() string
	// GetPriority 获取任务优先级
	GetPriority() int
	// GetCreatedAt 获取创建时间
	GetCreatedAt() time.Time
}

// TaskStatus 任务状态
type TaskStatus string

const (
	TaskStatusPending   TaskStatus = "pending"
	TaskStatusRunning   TaskStatus = "running"
	TaskStatusCompleted TaskStatus = "completed"
	TaskStatusFailed    TaskStatus = "failed"
	TaskStatusCancelled TaskStatus = "cancelled"
)

// TaskInfo 任务信息
type TaskInfo struct {
	ID        string
	Type      string
	Priority  int
	Status    TaskStatus
	CreatedAt time.Time
	StartedAt time.Time
	EndedAt   time.Time
	Error     string
}

// TaskQueue 任务队列接口
type TaskQueue interface {
	// AddTask 添加任务
	AddTask(task Task) error
	// GetTask 获取任务
	GetTask() (Task, error)
	// CompleteTask 完成任务
	CompleteTask(taskID string, err error) error
	// CancelTask 取消任务
	CancelTask(taskID string) error
	// GetTaskInfo 获取任务信息
	GetTaskInfo(taskID string) (*TaskInfo, error)
	// ListTasks 列出任务
	ListTasks() ([]*TaskInfo, error)
}

// PriorityQueue 基于优先级的任务队列
type PriorityQueue struct {
	tasks   []*taskItem
	mutex   sync.RWMutex
	cond    *sync.Cond
	taskMap map[string]*taskItem
}

// taskItem 任务项
type taskItem struct {
	task     Task
	info     *TaskInfo
	priority int
	index    int
}

// NewPriorityQueue 创建新的优先级队列
func NewPriorityQueue() *PriorityQueue {
	pq := &PriorityQueue{
		tasks:   make([]*taskItem, 0),
		taskMap: make(map[string]*taskItem),
	}
	pq.cond = sync.NewCond(&pq.mutex)
	return pq
}

// Len 队列长度
func (pq *PriorityQueue) Len() int {
	pq.mutex.RLock()
	defer pq.mutex.RUnlock()
	return len(pq.tasks)
}

// Less 比较两个任务的优先级
func (pq *PriorityQueue) Less(i, j int) bool {
	// 优先级数字越小优先级越高
	return pq.tasks[i].priority < pq.tasks[j].priority
}

// Swap 交换两个任务的位置
func (pq *PriorityQueue) Swap(i, j int) {
	pq.tasks[i], pq.tasks[j] = pq.tasks[j], pq.tasks[i]
	pq.tasks[i].index = i
	pq.tasks[j].index = j
}

// Push 添加任务
func (pq *PriorityQueue) Push(task Task) {
	pq.mutex.Lock()
	defer pq.mutex.Unlock()
	
	item := &taskItem{
		task:     task,
		priority: task.GetPriority(),
		index:    len(pq.tasks),
		info: &TaskInfo{
			ID:        task.GetID(),
			Type:      task.GetType(),
			Priority:  task.GetPriority(),
			Status:    TaskStatusPending,
			CreatedAt: task.GetCreatedAt(),
		},
	}
	
	pq.tasks = append(pq.tasks, item)
	pq.taskMap[task.GetID()] = item
	
	// 调整堆结构
	pq.up(len(pq.tasks) - 1)
	
	// 通知等待的消费者
	pq.cond.Signal()
}

// Pop 弹出最高优先级任务
func (pq *PriorityQueue) Pop() *taskItem {
	pq.mutex.Lock()
	defer pq.mutex.Unlock()
	
	for len(pq.tasks) == 0 {
		pq.cond.Wait()
	}
	
	item := pq.tasks[0]
	pq.remove(item.index)
	
	return item
}

// remove 移除任务
func (pq *PriorityQueue) remove(index int) {
	n := len(pq.tasks) - 1
	if index != n {
		pq.Swap(index, n)
		pq.down(index, n)
		pq.up(index)
	}
	
	item := pq.tasks[n]
	delete(pq.taskMap, item.info.ID)
	pq.tasks = pq.tasks[:n]
}

// up 向上调整堆结构
func (pq *PriorityQueue) up(j int) {
	for {
		i := (j - 1) / 2 // parent
		if i == j || !pq.Less(j, i) {
			break
		}
		pq.Swap(i, j)
		j = i
	}
}

// down 向下调整堆结构
func (pq *PriorityQueue) down(i, n int) {
	for {
		j1 := 2*i + 1
		if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
			break
		}
		j := j1 // left child
		if j2 := j1 + 1; j2 < n && pq.Less(j2, j1) {
			j = j2 // = 2*i + 2  // right child
		}
		if !pq.Less(j, i) {
			break
		}
		pq.Swap(i, j)
		i = j
	}
}

// AddTask 添加任务到队列
func (pq *PriorityQueue) AddTask(task Task) error {
	pq.Push(task)
	return nil
}

// GetTask 从队列获取任务
func (pq *PriorityQueue) GetTask() (Task, error) {
	item := pq.Pop()
	if item == nil {
		return nil, nil
	}
	
	// 更新任务状态
	item.info.Status = TaskStatusRunning
	item.info.StartedAt = time.Now()
	
	return item.task, nil
}

// CompleteTask 完成任务
func (pq *PriorityQueue) CompleteTask(taskID string, err error) error {
	pq.mutex.Lock()
	defer pq.mutex.Unlock()
	
	item, exists := pq.taskMap[taskID]
	if !exists {
		return nil
	}
	
	item.info.EndedAt = time.Now()
	if err != nil {
		item.info.Status = TaskStatusFailed
		item.info.Error = err.Error()
	} else {
		item.info.Status = TaskStatusCompleted
	}
	
	return nil
}

// CancelTask 取消任务
func (pq *PriorityQueue) CancelTask(taskID string) error {
	pq.mutex.Lock()
	defer pq.mutex.Unlock()
	
	item, exists := pq.taskMap[taskID]
	if !exists {
		return nil
	}
	
	item.info.Status = TaskStatusCancelled
	item.info.EndedAt = time.Now()
	
	// 从队列中移除
	pq.remove(item.index)
	
	return nil
}

// GetTaskInfo 获取任务信息
func (pq *PriorityQueue) GetTaskInfo(taskID string) (*TaskInfo, error) {
	pq.mutex.RLock()
	defer pq.mutex.RUnlock()
	
	item, exists := pq.taskMap[taskID]
	if !exists {
		return nil, nil
	}
	
	// 返回副本
	info := *item.info
	return &info, nil
}

// ListTasks 列出所有任务
func (pq *PriorityQueue) ListTasks() ([]*TaskInfo, error) {
	pq.mutex.RLock()
	defer pq.mutex.RUnlock()
	
	tasks := make([]*TaskInfo, 0, len(pq.taskMap))
	for _, item := range pq.taskMap {
		// 返回副本
		info := *item.info
		tasks = append(tasks, &info)
	}
	
	return tasks, nil
}

// Worker 工作协程接口
type Worker interface {
	// Start 启动工作协程
	Start(ctx context.Context)
	// Stop 停止工作协程
	Stop()
}

// TaskWorker 任务工作协程
type TaskWorker struct {
	id       string
	queue    TaskQueue
	ctx      context.Context
	cancel   context.CancelFunc
	wg       sync.WaitGroup
	onTaskStart func(task Task)
	onTaskEnd   func(task Task, err error)
}

// NewTaskWorker 创建新的任务工作协程
func NewTaskWorker(id string, queue TaskQueue) *TaskWorker {
	ctx, cancel := context.WithCancel(context.Background())
	return &TaskWorker{
		id:     id,
		queue:  queue,
		ctx:    ctx,
		cancel: cancel,
	}
}

// SetTaskCallbacks 设置任务回调
func (tw *TaskWorker) SetTaskCallbacks(onStart func(task Task), onEnd func(task Task, err error)) {
	tw.onTaskStart = onStart
	tw.onTaskEnd = onEnd
}

// Start 启动工作协程
func (tw *TaskWorker) Start() {
	tw.wg.Add(1)
	go func() {
		defer tw.wg.Done()
		tw.run()
	}()
}

// Stop 停止工作协程
func (tw *TaskWorker) Stop() {
	tw.cancel()
	tw.wg.Wait()
}

// run 运行工作协程
func (tw *TaskWorker) run() {
	for {
		select {
		case <-tw.ctx.Done():
			return
		default:
			// 获取任务
			task, err := tw.queue.GetTask()
			if err != nil {
				continue
			}
			
			if task == nil {
				continue
			}
			
			// 执行任务前回调
			if tw.onTaskStart != nil {
				tw.onTaskStart(task)
			}
			
			// 执行任务
			err = task.Execute(tw.ctx)
			
			// 完成任务
			tw.queue.CompleteTask(task.GetID(), err)
			
			// 执行任务后回调
			if tw.onTaskEnd != nil {
				tw.onTaskEnd(task, err)
			}
		}
	}
}

// TaskManager 任务管理器
type TaskManager struct {
	queue   TaskQueue
	workers []*TaskWorker
	mutex   sync.RWMutex
}

// NewTaskManager 创建新的任务管理器
func NewTaskManager(queue TaskQueue, workerCount int) *TaskManager {
	tm := &TaskManager{
		queue:   queue,
		workers: make([]*TaskWorker, 0, workerCount),
	}
	
	// 创建工作协程
	for i := 0; i < workerCount; i++ {
		worker := NewTaskWorker(fmt.Sprintf("worker-%d", i), queue)
		tm.workers = append(tm.workers, worker)
	}
	
	return tm
}

// Start 启动任务管理器
func (tm *TaskManager) Start() {
	tm.mutex.Lock()
	defer tm.mutex.Unlock()
	
	for _, worker := range tm.workers {
		worker.Start()
	}
}

// Stop 停止任务管理器
func (tm *TaskManager) Stop() {
	tm.mutex.Lock()
	defer tm.mutex.Unlock()
	
	for _, worker := range tm.workers {
		worker.Stop()
	}
}

// AddTask 添加任务
func (tm *TaskManager) AddTask(task Task) error {
	return tm.queue.AddTask(task)
}

// GetTaskInfo 获取任务信息
func (tm *TaskManager) GetTaskInfo(taskID string) (*TaskInfo, error) {
	return tm.queue.GetTaskInfo(taskID)
}

// ListTasks 列出所有任务
func (tm *TaskManager) ListTasks() ([]*TaskInfo, error) {
	return tm.queue.ListTasks()
}